package oneD.fem.structure;

import oneD.fem.algorithm.MyLinalg;
import inf.jlinalg.Array2DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.GeneralMatrixLSESolver;
import inf.jlinalg.ILSESolver;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import inf.jlinalg.JArpackSBMGESolver;
import inf.jlinalg.MatrixFormat;
import inf.jlinalg.QuadraticMatrixInfo;
import inf.jlinalg.SolveFailedException;
import inf.jlinalg.SymmetricBandMatrixInfo;

public class PrestressTruss extends Truss {

	private static double TOL = 1e-6;
	private static int MAX_ITERATION = 1000;

	public PrestressTruss() {

	}

	public void addElement(double rho, double e, double a, double preS, int n1,
			int n2) {
		this.elements_.add(new oneD.fem.element.nonlinear.PrestressTruss(rho, e, a, preS, this.nodes_
				.get(n1), this.nodes_.get(n2)));
	}

	@Override
	public void printStructure() {
		super.printStructure();
		System.out.println("Elements:");
		System.out.println("idx\tE\tA\tlength\tpre stress");
		for (int i = 0; i < this.getNumberOfElements(); i++) {
			System.out.print(i + "\t");
			this.elements_.get(i).print();
			System.out.println();
		}
	}

	public void assembleGeometricalStiffnessMatrix(IMatrix kgGlobal,
			boolean full) {
		for (int i = 0; i < this.getNumberOfElements(); i++) {
			IMatrix keg = this.elements_.get(i)
					.computeGeometricalStiffnessMatrix();
			int[] dofNumbers = this.elements_.get(i).getDOFNumbers();
			for (int p = 0; p < dofNumbers.length; p++)
				for (int q = full ? 0 : p; q < dofNumbers.length; q++)
					if (dofNumbers[p] != -1 && dofNumbers[q] != -1)
						kgGlobal.add(dofNumbers[p], dofNumbers[q],
								keg.get(p, q));
		}
		// System.out.println("System geometrical stiffness matrix:");
		// System.out.println(MatrixFormat.format(ktGlobal));
	}

	@Override
	public void assembleTangentStiffnessMatrix(IMatrix ktGlobal) {
		for (int i = 0; i < this.getNumberOfElements(); i++) {
			IMatrix ke = this.elements_.get(i).computeTangentStiffnessMatrix();
			int[] dofNumbers = this.elements_.get(i).getDOFNumbers();
			for (int p = 0; p < dofNumbers.length; p++)
				for (int q = 0; q < dofNumbers.length; q++)
					if (dofNumbers[p] != -1 && dofNumbers[q] != -1)
						ktGlobal.add(dofNumbers[p], dofNumbers[q], ke.get(p, q));
		}
		// System.out.println("System tangent stiffness matrix:");
		// System.out.println(MatrixFormat.format(ktGlobal));
	}

	@Override
	public void assembleInternalForceVector(IVector riGlobal) {
		for (int i = 0; i < this.getNumberOfElements(); i++) {
			// compute internal force vector of the element
			IVector rie = this.elements_.get(i).computeInternalForcesVector();
			// System.out.println("Element Internal forces vector when assemble");
			// System.out.println(MatrixFormat.format(rie));
			int[] dofNumbers = this.elements_.get(i).getDOFNumbers();
			// assemble the internal force vector to global internal force
			// vector
			for (int j = 0; j < dofNumbers.length; j++) {
				if (dofNumbers[j] != -1)
					riGlobal.add(dofNumbers[j], rie.get(j));
				// System.out.println(dofNumbers[j] + " " + rie.get(j));
			}
		}
	}

	@Override
	public void assembleLoadVector(IVector rGlobal) {
		for (int i = 0; i < this.getNumberOfNodes(); i++) {
			if (this.nodes_.get(i).getForce() != null) {
				int[] dofNumbers = this.nodes_.get(i).getDOFNumbers();
				for (int j = 0; j < 3; j++)
					if (dofNumbers[j] != -1)
						rGlobal.set(dofNumbers[j], this.nodes_.get(i)
								.getForce().getComponent(j));
			}
		}
	}

	@Override
	public void solve(IVector u0) {
		/*
		 * solve ri(u) = r using Newton Raphson scheme
		 */
		this.iterateNewtonRaphson(u0);
	}
	
	private void iterateNewtonRaphson(IVector u0) {
		int NEQ = this.enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
//		ILSESolver solver = new SpoolesLSESolver();
		
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		IMatrix ktGlobal = solver.getA();

		IVector rGlobal = new ArrayVector(NEQ);
		IVector riGlobal = new ArrayVector(NEQ);

		aInfo.setSize(NEQ);

		this.assembleLoadVector(rGlobal);

		double epsilon;
		// set displacement at each node to initial values
		this.setDisplacement(u0);
		int iteration = 1;
		do {
			solver.initialize();
			System.out.println("Analysis tracking step: " + (iteration++)
					+ "\n");
			BLAM.zero(ktGlobal);// clear ktGlobal to fill in new data
			BLAM.fill(riGlobal, 0.0);// clear global internal forces vector
			this.assembleTangentStiffnessMatrix(ktGlobal);
			this.assembleInternalForceVector(riGlobal);
//			System.out.println("Global tangent stiffness matrix:");
//			System.out.println(MatrixFormat.format(ktGlobal));
//			System.out.println("Global internal forces vector:");
//			System.out.println(MatrixFormat.format(riGlobal));
//			System.out.println();
//			System.out.println("Global external forces vector:");
//			System.out.println(MatrixFormat.format(rGlobal));
//			System.out.println();
			double[] b = null;
			try {
				// BLAM.add(-1, rGlobal, riGlobal);
				MyLinalg.add(-1, rGlobal, riGlobal);
				// BLAM.scale(-1, riGlobal);
				MyLinalg.scale(-1, riGlobal);
				b = MyLinalg.toArray(riGlobal);
				System.out.println("RHS vector:");
				MyLinalg.printVectorDouble(b);
				System.out.println();
				solver.solve(b);
			} catch (SolveFailedException e) {
				System.out.println("Solver failed: " + e.getMessage());
			}
			// update new displacement for each node
			this.updateDisplacement(b);
			System.out.println("Update vector:");
			MyLinalg.printVectorDouble(b);
			System.out.println();
			this.printDisplacement();
			// compute absolute error
			epsilon = MyLinalg.norm2(b);
			System.out.println("Absolute error = " + epsilon);
			System.out.println();
			System.out
					.println("/***********************************************************/");
			System.out.println();
		} while (epsilon > TOL && iteration < MAX_ITERATION);
		// this.assembleInternalForceVector(riGlobal);
		// System.out.println(MatrixFormat.format(rGlobal));
		// System.out.println(MatrixFormat.format(riGlobal));
		System.out.println("Iteration ended, step = " + iteration);
	}

	@Override
	public void directLoadControlInstabilityComputation(IVector u,
			double lambda, IVector phi) {
		int NEQ = this.enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		IMatrix K = solver.getA();
		IMatrix ktGlobal = new Array2DMatrix(NEQ, NEQ);
		IVector R = new ArrayVector(2 * NEQ + 1);
		IVector rGlobal = new ArrayVector(NEQ);
		IVector riGlobal = new ArrayVector(NEQ);
		IMatrix dKt_u_phi = new Array2DMatrix(NEQ, NEQ);

		aInfo.setSize(2 * NEQ + 1);

		this.assembleLoadVector(rGlobal);

		double epsilon;
		double eta = 1e-6; // parameter to numerically compute dKt_u
		// set displacement at each node to initial values
		int iteration = 1;
		do {
			solver.initialize();
			System.out.println("Analysis tracking step: " + (iteration++)
					+ "\n");

			BLAM.zero(ktGlobal);// clear ktGlobal to fill in new data
			BLAM.fill(riGlobal, 0.0);// clear global internal forces vector

			this.setDisplacement(u);
			this.assembleInternalForceVector(riGlobal);
			this.assembleTangentStiffnessMatrix(ktGlobal);

//			System.out.println("Tangent stiffness matrix");
//			System.out.println(MatrixFormat.format(ktGlobal));

			// compute dKt_u*phi
			for (int i = 0; i < NEQ; i++) {
				IVector ui = (IVector) u.clone();
				ui.add(i, eta);
				this.setDisplacement(ui);
				IMatrix ktNew = (IMatrix) ktGlobal.clone();
				BLAM.zero(ktNew);
				this.assembleTangentStiffnessMatrix(ktNew);
				BLAM.add(-1, ktGlobal, ktNew);
				BLAM.scale(1 / eta, ktNew);
				// reuse ui
				BLAM.multiply(1, BLAM.NO_TRANSPOSE, ktNew, phi, 0, ui);
				for (int j = 0; j < NEQ; j++)
					dKt_u_phi.set(j, i, ui.get(j));
			}
//			System.out.println("Derivative of Kt w.r.t u");
//			System.out.println(MatrixFormat.format(dKt_u_phi));

			// assemble K
			double normphi = BLAM.norm(BLAM.NORM_TWO, phi);
			BLAM.zero(K);
			for (int i = 0; i < NEQ; i++) {
				for (int j = 0; j < NEQ; j++) {
					K.set(i, j, ktGlobal.get(i, j));
					K.set(i + NEQ, j, dKt_u_phi.get(i, j));
					K.set(i + NEQ, j + NEQ + 1, ktGlobal.get(i, j));
				}
				K.set(i, NEQ, -rGlobal.get(i));
				K.set(2 * NEQ, i + NEQ + 1, phi.get(i) / normphi);
			}
//			System.out.println("Coupling matrix");
//			System.out.println(MatrixFormat.format(K));

			// assemble R
//			System.out.println("Internal forces vector");
//			System.out.println(MatrixFormat.format(riGlobal));
			// BLAM.add(-lambda, rGlobal, riGlobal);
			MyLinalg.add(-lambda, rGlobal, riGlobal);
			IVector kt_phi = (IVector) phi.clone();
			BLAM.multiply(1, BLAM.NO_TRANSPOSE, ktGlobal, phi, 0, kt_phi);
//			System.out.println("Kt*phi");
//			System.out.println(MatrixFormat.format(kt_phi));
			for (int i = 0; i < NEQ; i++) {
				R.set(i, -riGlobal.get(i));
				R.set(i + NEQ, -kt_phi.get(i));
			}
			R.set(2 * NEQ, 1 - normphi);
//			System.out.println("Coupling RHS vector");
//			System.out.println(MatrixFormat.format(R));

			double[] b = MyLinalg.toArray(R);
			try {
				solver.solve(b);
			} catch (SolveFailedException e) {
				System.out.println("Solver failed: " + e.getMessage());
			}

			// update new value
			for (int i = 0; i < NEQ; i++) {
				u.add(i, b[i]);
				phi.add(i, b[i + NEQ + 1]);
			}
			lambda += b[NEQ];

			epsilon = BLAM.norm(BLAM.NORM_TWO, R);
			System.out.println("Absolute error = " + epsilon);
			System.out.println();
			System.out
					.println("/***********************************************************/");
			System.out.println();
		} while (epsilon > TOL && iteration < MAX_ITERATION);
		System.out.println("Iteration ended, step = " + iteration);
		this.printDisplacement();
		System.out.println("Lambda = " + lambda);
		System.out.println("phi:");
		System.out.println(MatrixFormat.format(phi));
	}

	@Override
	public void secondOrderLoadControlInstabilityComputation(int nev, IVector u,
			IVector lambda) {
		int NEQ = this.enumerateDOFs();

		// IMatrix kLinGlobal = new Array2DMatrix(NEQ, NEQ);
		// IMatrix kGGlobal = new Array2DMatrix(NEQ, NEQ);
		//
		// this.assembleStiffnessMatrix(kLinGlobal, false);
		// this.setDisplacement(u);
		// this.assembleGeometricalStiffnessMatrix(kGGlobal, false);
		//
		// System.out.println("Linear stiffness matrix");
		// System.out.println(MatrixFormat.format(kLinGlobal));
		// System.out.println("Geometrical stiffness matrix");
		// System.out.println(MatrixFormat.format(kGGlobal));

		// solver initialization
		JArpackSBMGESolver solver = new JArpackSBMGESolver();
		SymmetricBandMatrixInfo aInfo = (SymmetricBandMatrixInfo) solver
				.getAInfo();
		SymmetricBandMatrixInfo bInfo = (SymmetricBandMatrixInfo) solver
				.getBInfo();

		// QuadraticMatrixInfo aInfo = solver.getAInfo();
		// QuadraticMatrixInfo bInfo = solver.getBInfo();

		IMatrix kLinGlobal = solver.getA();
		IMatrix kGGlobal = solver.getB();

		aInfo.setSize(NEQ);
		bInfo.setSize(NEQ);

		// aInfo.setSymmetric(true);
		// bInfo.setSymmetric(true);

		// aInfo.setBandwidth(NEQ - 1);
		// bInfo.setBandwidth(NEQ - 1);
		aInfo.setBandwidth(2 * NEQ);
		bInfo.setBandwidth(2 * NEQ);

		solver.setNev(nev);
		solver.setMode(JArpackSBMGESolver.MODE_BUCKLING);

		solver.initialize();

		this.assembleStiffnessMatrix(kLinGlobal, false);
		this.setDisplacement(u);
		this.assembleGeometricalStiffnessMatrix(kGGlobal, false);

		System.out.println("Linear stiffness matrix");
		System.out.println(MatrixFormat.format(kLinGlobal));
		System.out.println("Geometrical stiffness matrix");
		System.out.println(MatrixFormat.format(kGGlobal));

		try {
			solver.solve();
		} catch (SolveFailedException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

		lambda = (IVector) solver.getEigenvalues();
		MyLinalg.scale(-1, lambda);
		System.out.println("Instability load factors:");
		System.out.println(MatrixFormat.format(lambda));
	}
}
